import os
import torch
import torch.utils.data
import torchvision
import numpy as np
import pandas as pd
from data.apple_dataset import AppleDataset
import utility.utils as utils
import utility.transforms as T
import datetime
import time
import matplotlib.pyplot as plt
import torch
import torch.utils.data
from PIL import Image
import torchvision
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from torchvision.models.detection.mask_rcnn import MaskRCNNPredictor

from data.apple_dataset import AppleDataset
from utility.engine import train_one_epoch, evaluate

import utility.utils as utils
import utility.transforms as T

def get_transform(train):
    transforms = []
    transforms.append(T.ToTensor())
    if train:
        transforms.append(T.RandomHorizontalFlip(0.5))
    return T.Compose(transforms)


def save_mask_as_image(mask, output_path, img_name, mask_id,data_path):
    if mask.dtype != np.float32:
        mask = mask.astype(np.float32)
    if mask.min() < 0 or mask.max() > 1:
        mask = (mask - mask.min()) / (mask.max() - mask.min())
    print(mask.shape)
    if len(mask.shape) == 4:
        mask = mask[0, 0, :, :]
    elif len(mask.shape) == 3:
        mask = mask[0, :, :]

    mask = torch.from_numpy(mask)
    os.makedirs(output_path, exist_ok=True)

    mask_path = os.path.join(output_path,"mask", f"{img_name}")
    try:
        # torchvision.utils.save_image(mask, mask_path)
        # print(f"Saved mask to {mask_path}")

        # original_img_path = os.path.join(data_path, 'images', f"{img_name}")
        # compare_path = os.path.join(output_path, "compared")
        # os.makedirs(compare_path, exist_ok=True)
        #
        # if os.path.exists(original_img_path):
        #     # 打开原始图像
        #     original_img = Image.open(original_img_path)
        #     print(original_img.size)
        #     # 将mask转换为PIL图像
        #     mask_img = Image.fromarray((mask.numpy().squeeze() * 255).astype(np.uint8))
        #
        #     # 创建红色半透明层
        #     red_layer = Image.new('RGBA', original_img.size, (255, 0, 0, 76))  # 76表示30%透明度
        #
        #     # 创建最终图像
        #     final_img = original_img.convert('RGBA')
        #     final_img.paste(red_layer, (0, 0), mask_img)
        #
        #     # 保存结果
        #     compare_img_path = os.path.join(compare_path, img_name)
        #     final_img.save(compare_img_path)
        #     print(f"Saved comparison to {compare_img_path}")
        original_img_path = os.path.join(data_path, 'images', f"{img_name}")
        compare_path = os.path.join(output_path, "compared")
        os.makedirs(compare_path, exist_ok=True)
        print(original_img_path)
        if os.path.exists(original_img_path):
            # 读取原图
            original_img = plt.imread(original_img_path)

            # 创建固定尺寸的画布（720x1280）
            fig, ax = plt.subplots(figsize=(7.2, 12.8), dpi=100)
            plt.subplots_adjust(left=0, right=1, top=1, bottom=0)
            # 保证一定要输出全部占满画布

            # 显示原图并调整尺寸
            ax.imshow(original_img, extent=[0, 720, 0, 1280])  # 强制设置为720x1280
            ax.axis('off')
            # plt.savefig('temp2.png', pad_inches=0, dpi=100)
            # 显示掩码叠加效果
            ax.imshow(mask.numpy().squeeze(), cmap='Reds', alpha=0.3, extent=[0, 720, 0, 1280])

            # 保存图片
            compare_img_path = os.path.join(compare_path, img_name)
            plt.savefig(compare_img_path, pad_inches=0, dpi=100)
            plt.close()
            print(f"Saved comparison to {compare_img_path}")
    except Exception as e:
        print(f"Error saving mask: {e}")
def predict_with_gui(data_path, weight_file, model_type='mrcnn', device='cuda'):
    """
    GUI版本的预测函数
    参数:
        data_path: 测试数据路径
        weight_file: 模型权重文件路径
        model_type: 模型类型 ('mrcnn' 或 'frcnn')
        device: 计算设备 ('cuda' 或 'cpu')
    返回:
        results: 包含预测结果的列表，每个元素是一个字典，包含图像名称和预测框信息
    """
    # 初始化设备
    device = torch.device(device if torch.cuda.is_available() else 'cpu')

    # 加载模型
    print("Loading model...")
    model = torch.load(weight_file, map_location=device)
    model.to(device)
    model.eval()

    # 创建数据加载器
    print("Creating data loader...")
    dataset_test = AppleDataset(data_path, get_transform(train=False))
    data_loader_test = torch.utils.data.DataLoader(
        dataset_test, batch_size=1, shuffle=False,
        num_workers=1, collate_fn=utils.collate_fn
    )

    # 存储预测结果
    results = []

    # 设置 mask 输出路径
    output_mask_path = "./predict2output"
    os.makedirs(output_mask_path, exist_ok=True)

    # 开始预测
    print("Starting prediction...")
    with torch.no_grad():
        for image, targets in data_loader_test:
            # 移动数据到设备
            image = list(img.to(device) for img in image)

            # 模型预测
            outputs = model(image)

            # 处理每个预测结果
            for ii, output in enumerate(outputs):
                img_id = targets[ii]['image_id']
                img_name = dataset_test.get_img_name(img_id)

                # 获取预测数据时添加内存优化
                boxes = output['boxes'].detach().cpu().numpy()
                scores = output['scores'].detach().cpu().numpy()

                # 添加置信度过滤 (阈值设为0.5)
                keep_idx = scores > 0.5
                boxes = boxes[keep_idx]
                scores = scores[keep_idx]

                # 分批处理mask防止内存溢出
                combined_mask = None
                for mask_idx in range(output['masks'].shape[0]):
                    if keep_idx[mask_idx]:  # 只处理通过过滤的mask
                        # 单mask处理
                        mask = output['masks'][mask_idx, 0].detach().cpu().numpy()
                        binary_mask = (mask > 0.5).astype(np.uint8) * 255

                        # 合并mask
                        if combined_mask is None:
                            combined_mask = binary_mask
                        else:
                            combined_mask = np.maximum(combined_mask, binary_mask)

                        # 及时释放内存
                        del mask, binary_mask

                # 保存合并结果
                if combined_mask is not None:
                    save_mask_as_image(combined_mask, output_mask_path, img_name, "combined",data_path)
                    del combined_mask

                # 存储精简后的结果
                result = {
                    'image_name': img_name,
                    'boxes': boxes,
                    'scores': scores
                }
                results.append(result)

                # 手动触发垃圾回收
                if ii % 10 == 0:
                    torch.cuda.empty_cache()

                print(f"Predicted on image: {img_name}")

                # 保存 mask 到图片
                # for mask_idx, mask in enumerate(masks):
                #     # 取第一个 mask (假设只预测一个实例)
                #     mask = mask[0, :, :]  # shape: [height, width]
                #     save_mask_as_image(mask, output_mask_path, img_name, mask_idx)
                # print(masks);
                # a=sum(sum(sum(sum(masks))));
                # print(f"Predicted on image: {img_name}")

    print("Prediction completed!")
    return results
def main():
    # 固定参数设置
    args = {
        'data_path': r"D:/Jupyter/pytorch/d2l-zh/pytorch/appppppp/detection2/test",
        'output_file': "./output/1.txt",
        'weight_file': r"E:/download/MinneApple-master/MinneApple-master/MinneApple-master/myapple101.pth",
        'device': "cuda",
        'model_type': 'mrcnn'  # 直接指定模型类型
    }
    print("Initializing prediction...")
    # 执行预测
    results = predict_with_gui(
        data_path=args['data_path'],
        weight_file=args['weight_file'],
        model_type=args['model_type'],
        device=args['device']
    )

    # 保存结果到文件
    with open(args['output_file'], 'w') as f:
        for result in results:
            boxes = result['boxes'].astype(int)
            scores = result['scores']
            for box, score in zip(boxes, scores):
                line = f"{result['image_name']},{box[0]},{box[1]},{box[2]},{box[3]},{score:.4f}\n"
                f.write(line)
    print(f"预测结果已保存至 {args['output_file']}")
    df_list = []
    for result in results:
        for box, score in zip(result['boxes'], result['scores']):
            df_list.append([
                result['image_name'],
                *box.astype(int).tolist(),
                round(float(score), 4)
            ])

    pred_df = pd.DataFrame(df_list,
                           columns=['image_name', 'x1', 'y1', 'x2', 'y2', 'confidence'])

    # 保存结果
    pred_df.to_csv(args['output_file'], index=False)
    print(f"Saved predictions to {args['output_file']}")

if __name__ == "__main__":
    main()